import os
import Master
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, WebSocketException


# load_dotenv()

app = FastAPI()
master = Master.Master()

@app.get("/")
def root():
    return {"message": os.getenv("DEEPSEEK_API_KEY")}

@app.post("/chat")
def chat(query:str):
    return {"message": master.run(query)}

@app.post("/add_urls")
def add_urls():
    return {"message": "add_urls"}

@app.post("/add_pdfs")
def add_pdfs():
    return {"message": "add_pdfs"}

@app.post("/add_texts")
def add_texts():
    return {"message": "add_texts"}

@app.websocket("/ws")
async def websocket_endpoint(websocket : WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            result = master.run( data)
            if result:
                await websocket.send_text(result["output"])
            # await websocket.send_text(f"Message text was: {data}")
    except WebSocketException:
        await websocket.close()


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="localhost", port=8000)